{"cells":[{"cell_type":"markdown","source":["Copyright 2022 Google LLC. SPDX-License-Identifier: Apache-2.0\n","\n","Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n","\n","https://www.apache.org/licenses/LICENSE-2.0\n","\n","Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."],"metadata":{"id":"WwHXse056QPQ"}},{"cell_type":"markdown","source":["# Code as Policies Tabletop Manipulation Interactive Demo\n","\n","This notebook is a part of the open-source code release associated with the paper:\n","\n","[Code as Policies: Language Model Programs for Embodied Control](https://code-as-policies.github.io/)\n","\n","This notebook gives an interactive demo for the simulated tabletop manipulation domain, where users can type in free-form instructions for the robot to execute, as seen in the paper section IV.D\n","\n","**Instructions:** \n","\n","1) Obtain an OpenAI API Key in the link below and set it in the next cell:\n","https://openai.com/blog/openai-api/\n","\n","2) Gain Codex access by joining the waitlist here:\n","https://openai.com/blog/openai-codex/\n","\n","Once you have Codex access you can use `code-davinci-002` as the `model_name`. Using the GPT-3 model (`text-dainvci-002`) is also ok, but performance won't be as good (there will be more code logic errors).\n","\n","3) Run all cells and input commands under \"Interactive Demo\" at the bottom."],"metadata":{"id":"nWKW9oUa6SFJ"}},{"cell_type":"code","source":["openai_api_key = 'YOUR KEY HERE'\n","model_name = 'code-davinci-002' # 'text-davinci-002'"],"metadata":{"id":"DIufITOa6k-i"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"jOvZPzEEm6tv"},"source":["# Setup"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"IqLFRuhdmpB5"},"outputs":[],"source":["!pip install numpy scipy shapely astunparse pygments openai > /dev/null 2>&1\n","!pip install imageio==2.4.1 imageio-ffmpeg pybullet moviepy\n","\n","import os\n","import pybullet\n","import pybullet_data\n","import numpy as np\n","import threading\n","import copy\n","import openai\n","import cv2\n","from google.colab.patches import cv2_imshow\n","from moviepy.editor import ImageSequenceClip\n","\n","# imports for LMPs\n","import shapely\n","import ast\n","import astunparse\n","from time import sleep\n","from shapely.geometry import *\n","from shapely.affinity import *\n","from openai.error import RateLimitError, APIConnectionError\n","from pygments import highlight\n","from pygments.lexers import PythonLexer\n","from pygments.formatters import TerminalFormatter\n","\n","openai.api_key = openai_api_key\n","\n","# Download PyBullet assets.\n","if not os.path.exists('ur5e/ur5e.urdf'):\n","  !gdown --id 1Cc_fDSBL6QiDvNT4dpfAEbhbALSVoWcc\n","  !gdown --id 1yOMEm-Zp_DL3nItG9RozPeJAmeOldekX\n","  !gdown --id 1GsqNLhEl9dd4Mc3BM0dX3MibOI1FVWNM\n","  !unzip ur5e.zip\n","  !unzip robotiq_2f_85.zip\n","  !unzip bowl.zip\n","\n","!nvidia-smi  # Show useful GPU info."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"hjibCwPlncWZ"},"outputs":[],"source":["class LMP:\n","\n","    def __init__(self, name, cfg, lmp_fgen, fixed_vars, variable_vars):\n","        self._name = name\n","        self._cfg = cfg\n","\n","        self._base_prompt = self._cfg['prompt_text']\n","\n","        self._stop_tokens = list(self._cfg['stop'])\n","\n","        self._lmp_fgen = lmp_fgen\n","\n","        self._fixed_vars = fixed_vars\n","        self._variable_vars = variable_vars\n","        self.exec_hist = ''\n","\n","    def clear_exec_hist(self):\n","        self.exec_hist = ''\n","\n","    def build_prompt(self, query, context=''):\n","        if len(self._variable_vars) > 0:\n","            variable_vars_imports_str = f\"from utils import {', '.join(self._variable_vars.keys())}\"\n","        else:\n","            variable_vars_imports_str = ''\n","        prompt = self._base_prompt.replace('{variable_vars_imports}', variable_vars_imports_str)\n","\n","        if self._cfg['maintain_session']:\n","            prompt += f'\\n{self.exec_hist}'\n","\n","        if context != '':\n","            prompt += f'\\n{context}'\n","\n","        use_query = f'{self._cfg[\"query_prefix\"]}{query}{self._cfg[\"query_suffix\"]}'\n","        prompt += f'\\n{use_query}'\n","\n","        return prompt, use_query\n","\n","    def __call__(self, query, context='', **kwargs):\n","        prompt, use_query = self.build_prompt(query, context=context)\n","\n","        while True:\n","            try:\n","                code_str = openai.Completion.create(\n","                    prompt=prompt,\n","                    stop=self._stop_tokens,\n","                    temperature=self._cfg['temperature'],\n","                    engine=self._cfg['engine'],\n","                    max_tokens=self._cfg['max_tokens']\n","                )['choices'][0]['text'].strip()\n","                break\n","            except (RateLimitError, APIConnectionError) as e:\n","                print(f'OpenAI API got err {e}')\n","                print('Retrying after 10s.')\n","                sleep(10)\n","\n","        if self._cfg['include_context'] and context != '':\n","            to_exec = f'{context}\\n{code_str}'\n","            to_log = f'{context}\\n{use_query}\\n{code_str}'\n","        else:\n","            to_exec = code_str\n","            to_log = f'{use_query}\\n{to_exec}'\n","\n","        to_log_pretty = highlight(to_log, PythonLexer(), TerminalFormatter())\n","        print(f'LMP {self._name} exec:\\n\\n{to_log_pretty}\\n')\n","\n","        new_fs = self._lmp_fgen.create_new_fs_from_code(code_str)\n","        self._variable_vars.update(new_fs)\n","\n","        gvars = merge_dicts([self._fixed_vars, self._variable_vars])\n","        lvars = kwargs\n","\n","        if not self._cfg['debug_mode']:\n","            exec_safe(to_exec, gvars, lvars)\n","\n","        self.exec_hist += f'\\n{to_exec}'\n","\n","        if self._cfg['maintain_session']:\n","            self._variable_vars.update(lvars)\n","\n","        if self._cfg['has_return']:\n","            return lvars[self._cfg['return_val_name']]\n","\n","\n","class LMPFGen:\n","\n","    def __init__(self, cfg, fixed_vars, variable_vars):\n","        self._cfg = cfg\n","\n","        self._stop_tokens = list(self._cfg['stop'])\n","        self._fixed_vars = fixed_vars\n","        self._variable_vars = variable_vars\n","\n","        self._base_prompt = self._cfg['prompt_text']\n","\n","    def create_f_from_sig(self, f_name, f_sig, other_vars=None, fix_bugs=False, return_src=False):\n","        print(f'Creating function: {f_sig}')\n","\n","        use_query = f'{self._cfg[\"query_prefix\"]}{f_sig}{self._cfg[\"query_suffix\"]}'\n","        prompt = f'{self._base_prompt}\\n{use_query}'\n","\n","        while True:\n","            try:\n","                f_src = openai.Completion.create(\n","                    prompt=prompt, \n","                    stop=self._stop_tokens,\n","                    temperature=self._cfg['temperature'],\n","                    engine=self._cfg['engine'],\n","                    max_tokens=self._cfg['max_tokens']\n","                )['choices'][0]['text'].strip()\n","                break\n","            except (RateLimitError, APIConnectionError) as e:\n","                print(f'OpenAI API got err {e}')\n","                print('Retrying after 10s.')\n","                sleep(10)\n","\n","        if fix_bugs:\n","            f_src = openai.Edit.create(\n","                model='code-davinci-edit-001',\n","                input='# ' + f_src,\n","                temperature=0,\n","                instruction='Fix the bug if there is one. Improve readability. Keep same inputs and outputs. Only small changes. No comments.',\n","            )['choices'][0]['text'].strip()\n","\n","        if other_vars is None:\n","            other_vars = {}\n","        gvars = merge_dicts([self._fixed_vars, self._variable_vars, other_vars])\n","        lvars = {}\n","        \n","        exec_safe(f_src, gvars, lvars)\n","\n","        f = lvars[f_name]\n","\n","        to_print = highlight(f'{use_query}\\n{f_src}', PythonLexer(), TerminalFormatter())\n","        print(f'LMP FGEN created:\\n\\n{to_print}\\n')\n","\n","        if return_src:\n","            return f, f_src\n","        return f\n","\n","    def create_new_fs_from_code(self, code_str, other_vars=None, fix_bugs=False, return_src=False):\n","        fs, f_assigns = {}, {}\n","        f_parser = FunctionParser(fs, f_assigns)\n","        f_parser.visit(ast.parse(code_str))\n","        for f_name, f_assign in f_assigns.items():\n","            if f_name in fs:\n","                fs[f_name] = f_assign\n","\n","        if other_vars is None:\n","            other_vars = {}\n","\n","        new_fs = {}\n","        srcs = {}\n","        for f_name, f_sig in fs.items():\n","            all_vars = merge_dicts([self._fixed_vars, self._variable_vars, new_fs, other_vars])\n","            if not var_exists(f_name, all_vars):\n","                f, f_src = self.create_f_from_sig(f_name, f_sig, new_fs, fix_bugs=fix_bugs, return_src=True)\n","\n","                # recursively define child_fs in the function body if needed\n","                f_def_body = astunparse.unparse(ast.parse(f_src).body[0].body)\n","                child_fs, child_f_srcs = self.create_new_fs_from_code(\n","                    f_def_body, other_vars=all_vars, fix_bugs=fix_bugs, return_src=True\n","                )\n","\n","                if len(child_fs) > 0:\n","                    new_fs.update(child_fs)\n","                    srcs.update(child_f_srcs)\n","\n","                    # redefine parent f so newly created child_fs are in scope\n","                    gvars = merge_dicts([self._fixed_vars, self._variable_vars, new_fs, other_vars])\n","                    lvars = {}\n","                    \n","                    exec_safe(f_src, gvars, lvars)\n","                    \n","                    f = lvars[f_name]\n","\n","                new_fs[f_name], srcs[f_name] = f, f_src\n","\n","        if return_src:\n","            return new_fs, srcs\n","        return new_fs\n","\n","\n","class FunctionParser(ast.NodeTransformer):\n","\n","    def __init__(self, fs, f_assigns):\n","      super().__init__()\n","      self._fs = fs\n","      self._f_assigns = f_assigns\n","\n","    def visit_Call(self, node):\n","        self.generic_visit(node)\n","        if isinstance(node.func, ast.Name):\n","            f_sig = astunparse.unparse(node).strip()\n","            f_name = astunparse.unparse(node.func).strip()\n","            self._fs[f_name] = f_sig\n","        return node\n","\n","    def visit_Assign(self, node):\n","        self.generic_visit(node)\n","        if isinstance(node.value, ast.Call):\n","            assign_str = astunparse.unparse(node).strip()\n","            f_name = astunparse.unparse(node.value.func).strip()\n","            self._f_assigns[f_name] = assign_str\n","        return node\n","\n","\n","def var_exists(name, all_vars):\n","    try:\n","        eval(name, all_vars)\n","    except:\n","        exists = False\n","    else:\n","        exists = True\n","    return exists\n","\n","\n","def merge_dicts(dicts):\n","    return {\n","        k : v \n","        for d in dicts\n","        for k, v in d.items()\n","    }\n","    \n","\n","def exec_safe(code_str, gvars=None, lvars=None):\n","    banned_phrases = ['import', '__']\n","    for phrase in banned_phrases:\n","        assert phrase not in code_str\n","  \n","    if gvars is None:\n","        gvars = {}\n","    if lvars is None:\n","        lvars = {}\n","    empty_fn = lambda *args, **kwargs: None\n","    custom_gvars = merge_dicts([\n","        gvars,\n","        {'exec': empty_fn, 'eval': empty_fn}\n","    ])\n","    exec(code_str, custom_gvars, lvars)"]},{"cell_type":"markdown","metadata":{"id":"6YErUBSo3RT2"},"source":["## **Table Top Sim Env**\n","Define PyBullet-based environment with a UR5e and 2-finger gripper"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"D6gTAdeX39yd"},"outputs":[],"source":["# # Global constants: pick and place objects, colors, workspace bounds\n","COLORS = {\n","    'blue':   (78/255,  121/255, 167/255, 255/255),\n","    'red':    (255/255,  87/255,  89/255, 255/255),\n","    'green':  (89/255,  169/255,  79/255, 255/255),\n","    'orange': (242/255, 142/255,  43/255, 255/255),\n","    'yellow': (237/255, 201/255,  72/255, 255/255),\n","    'purple': (176/255, 122/255, 161/255, 255/255),\n","    'pink':   (255/255, 157/255, 167/255, 255/255),\n","    'cyan':   (118/255, 183/255, 178/255, 255/255),\n","    'brown':  (156/255, 117/255,  95/255, 255/255),\n","    'gray':   (186/255, 176/255, 172/255, 255/255),\n","}\n","\n","CORNER_POS = {\n","  'top left corner':     (-0.3 + 0.05, -0.2 - 0.05, 0),\n","  'top side':            (0,           -0.2 - 0.05, 0),\n","  'top right corner':    (0.3 - 0.05,  -0.2 - 0.05, 0),\n","  'left side':           (-0.3 + 0.05, -0.5,        0),\n","  'middle':              (0,           -0.5,        0),\n","  'right side':          (0.3 - 0.05,  -0.5,        0),\n","  'bottom left corner':  (-0.3 + 0.05, -0.8 + 0.05, 0),\n","  'bottom side':         (0,           -0.8 + 0.05, 0),\n","  'bottom right corner': (0.3 - 0.05,  -0.8 + 0.05, 0),\n","}\n","\n","ALL_BLOCKS = ['blue block', 'red block', 'green block', 'orange block', 'yellow block', 'purple block', 'pink block', 'cyan block', 'brown block', 'gray block']\n","ALL_BOWLS = ['blue bowl', 'red bowl', 'green bowl', 'orange bowl', 'yellow bowl', 'purple bowl', 'pink bowl', 'cyan bowl', 'brown bowl', 'gray bowl']\n","\n","PIXEL_SIZE = 0.00267857\n","BOUNDS = np.float32([[-0.3, 0.3], [-0.8, -0.2], [0, 0.15]])  # X Y Z"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"0bj91lBq6WeK"},"outputs":[],"source":["# Gripper (Robotiq 2F85) code\n","\n","class Robotiq2F85:\n","  \"\"\"Gripper handling for Robotiq 2F85.\"\"\"\n","\n","  def __init__(self, robot, tool):\n","    self.robot = robot\n","    self.tool = tool\n","    pos = [0.1339999999999999, -0.49199999999872496, 0.5]\n","    rot = pybullet.getQuaternionFromEuler([np.pi, 0, np.pi])\n","    urdf = 'robotiq_2f_85/robotiq_2f_85.urdf'\n","    self.body = pybullet.loadURDF(urdf, pos, rot)\n","    self.n_joints = pybullet.getNumJoints(self.body)\n","    self.activated = False\n","\n","    # Connect gripper base to robot tool.\n","    pybullet.createConstraint(self.robot, tool, self.body, 0, jointType=pybullet.JOINT_FIXED, jointAxis=[0, 0, 0], parentFramePosition=[0, 0, 0], childFramePosition=[0, 0, -0.07], childFrameOrientation=pybullet.getQuaternionFromEuler([0, 0, np.pi / 2]))\n","\n","    # Set friction coefficients for gripper fingers.\n","    for i in range(pybullet.getNumJoints(self.body)):\n","      pybullet.changeDynamics(self.body, i, lateralFriction=10.0, spinningFriction=1.0, rollingFriction=1.0, frictionAnchor=True)\n","\n","    # Start thread to handle additional gripper constraints.\n","    self.motor_joint = 1\n","    self.constraints_thread = threading.Thread(target=self.step)\n","    self.constraints_thread.daemon = True\n","    self.constraints_thread.start()\n","\n","  # Control joint positions by enforcing hard contraints on gripper behavior.\n","  # Set one joint as the open/close motor joint (other joints should mimic).\n","  def step(self):\n","    while True:\n","      try:\n","        currj = [pybullet.getJointState(self.body, i)[0] for i in range(self.n_joints)]\n","        indj = [6, 3, 8, 5, 10]\n","        targj = [currj[1], -currj[1], -currj[1], currj[1], currj[1]]\n","        pybullet.setJointMotorControlArray(self.body, indj, pybullet.POSITION_CONTROL, targj, positionGains=np.ones(5))\n","      except:\n","        return\n","      sleep(0.001)\n","\n","  # Close gripper fingers.\n","  def activate(self):\n","    pybullet.setJointMotorControl2(self.body, self.motor_joint, pybullet.VELOCITY_CONTROL, targetVelocity=1, force=10)\n","    self.activated = True\n","\n","  # Open gripper fingers.\n","  def release(self):\n","    pybullet.setJointMotorControl2(self.body, self.motor_joint, pybullet.VELOCITY_CONTROL, targetVelocity=-1, force=10)\n","    self.activated = False\n","\n","  # If activated and object in gripper: check object contact.\n","  # If activated and nothing in gripper: check gripper contact.\n","  # If released: check proximity to surface (disabled).\n","  def detect_contact(self):\n","    obj, _, ray_frac = self.check_proximity()\n","    if self.activated:\n","      empty = self.grasp_width() < 0.01\n","      cbody = self.body if empty else obj\n","      if obj == self.body or obj == 0:\n","        return False\n","      return self.external_contact(cbody)\n","  #   else:\n","  #     return ray_frac < 0.14 or self.external_contact()\n","\n","  # Return if body is in contact with something other than gripper\n","  def external_contact(self, body=None):\n","    if body is None:\n","      body = self.body\n","    pts = pybullet.getContactPoints(bodyA=body)\n","    pts = [pt for pt in pts if pt[2] != self.body]\n","    return len(pts) > 0  # pylint: disable=g-explicit-length-test\n","\n","  def check_grasp(self):\n","    while self.moving():\n","      sleep(0.001)\n","    success = self.grasp_width() > 0.01\n","    return success\n","\n","  def grasp_width(self):\n","    lpad = np.array(pybullet.getLinkState(self.body, 4)[0])\n","    rpad = np.array(pybullet.getLinkState(self.body, 9)[0])\n","    dist = np.linalg.norm(lpad - rpad) - 0.047813\n","    return dist\n","\n","  def check_proximity(self):\n","    ee_pos = np.array(pybullet.getLinkState(self.robot, self.tool)[0])\n","    tool_pos = np.array(pybullet.getLinkState(self.body, 0)[0])\n","    vec = (tool_pos - ee_pos) / np.linalg.norm((tool_pos - ee_pos))\n","    ee_targ = ee_pos + vec\n","    ray_data = pybullet.rayTest(ee_pos, ee_targ)[0]\n","    obj, link, ray_frac = ray_data[0], ray_data[1], ray_data[2]\n","    return obj, link, ray_frac"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"eWUtvw1qZi_5"},"outputs":[],"source":["# Gym-style environment code\n","\n","class PickPlaceEnv():\n","\n","  def __init__(self, render=False, high_res=False, high_frame_rate=False):\n","    self.dt = 1/480\n","    self.sim_step = 0\n","\n","    # Configure and start PyBullet.\n","    # python3 -m pybullet_utils.runServer\n","    # pybullet.connect(pybullet.SHARED_MEMORY)  # pybullet.GUI for local GUI.\n","    pybullet.connect(pybullet.DIRECT)  # pybullet.GUI for local GUI.\n","    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_GUI, 0)\n","    pybullet.setPhysicsEngineParameter(enableFileCaching=0)\n","    assets_path = os.path.dirname(os.path.abspath(\"\"))\n","    pybullet.setAdditionalSearchPath(assets_path)\n","    pybullet.setAdditionalSearchPath(pybullet_data.getDataPath())\n","    pybullet.setTimeStep(self.dt)\n","\n","    self.home_joints = (np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 3 * np.pi / 2, 0)  # Joint angles: (J0, J1, J2, J3, J4, J5).\n","    self.home_ee_euler = (np.pi, 0, np.pi)  # (RX, RY, RZ) rotation in Euler angles.\n","    self.ee_link_id = 9  # Link ID of UR5 end effector.\n","    self.tip_link_id = 10  # Link ID of gripper finger tips.\n","    self.gripper = None\n","\n","    self.render = render\n","    self.high_res = high_res\n","    self.high_frame_rate = high_frame_rate\n","\n","  def reset(self, object_list):\n","    pybullet.resetSimulation(pybullet.RESET_USE_DEFORMABLE_WORLD)\n","    pybullet.setGravity(0, 0, -9.8)\n","    self.cache_video = []\n","\n","    # Temporarily disable rendering to load URDFs faster.\n","    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 0)\n","\n","    # Add robot.\n","    pybullet.loadURDF(\"plane.urdf\", [0, 0, -0.001])\n","    self.robot_id = pybullet.loadURDF(\"ur5e/ur5e.urdf\", [0, 0, 0], flags=pybullet.URDF_USE_MATERIAL_COLORS_FROM_MTL)\n","    self.ghost_id = pybullet.loadURDF(\"ur5e/ur5e.urdf\", [0, 0, -10])  # For forward kinematics.\n","    self.joint_ids = [pybullet.getJointInfo(self.robot_id, i) for i in range(pybullet.getNumJoints(self.robot_id))]\n","    self.joint_ids = [j[0] for j in self.joint_ids if j[2] == pybullet.JOINT_REVOLUTE]\n","\n","    # Move robot to home configuration.\n","    for i in range(len(self.joint_ids)):\n","      pybullet.resetJointState(self.robot_id, self.joint_ids[i], self.home_joints[i])\n","\n","    # Add gripper.\n","    if self.gripper is not None:\n","      while self.gripper.constraints_thread.is_alive():\n","        self.constraints_thread_active = False\n","    self.gripper = Robotiq2F85(self.robot_id, self.ee_link_id)\n","    self.gripper.release()\n","\n","    # Add workspace.\n","    plane_shape = pybullet.createCollisionShape(pybullet.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])\n","    plane_visual = pybullet.createVisualShape(pybullet.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])\n","    plane_id = pybullet.createMultiBody(0, plane_shape, plane_visual, basePosition=[0, -0.5, 0])\n","    pybullet.changeVisualShape(plane_id, -1, rgbaColor=[0.2, 0.2, 0.2, 1.0])\n","\n","    # Load objects according to config.\n","    self.object_list = object_list\n","    self.obj_name_to_id = {}\n","    obj_xyz = np.zeros((0, 3))\n","    for obj_name in object_list:\n","      if ('block' in obj_name) or ('bowl' in obj_name):\n","\n","        # Get random position 15cm+ from other objects.\n","        while True:\n","          rand_x = np.random.uniform(BOUNDS[0, 0] + 0.1, BOUNDS[0, 1] - 0.1)\n","          rand_y = np.random.uniform(BOUNDS[1, 0] + 0.1, BOUNDS[1, 1] - 0.1)\n","          rand_xyz = np.float32([rand_x, rand_y, 0.03]).reshape(1, 3)\n","          if len(obj_xyz) == 0:\n","            obj_xyz = np.concatenate((obj_xyz, rand_xyz), axis=0)\n","            break\n","          else:\n","            nn_dist = np.min(np.linalg.norm(obj_xyz - rand_xyz, axis=1)).squeeze()\n","            if nn_dist > 0.15:\n","              obj_xyz = np.concatenate((obj_xyz, rand_xyz), axis=0)\n","              break\n","        \n","        object_color = COLORS[obj_name.split(' ')[0]]\n","        object_type = obj_name.split(' ')[1]\n","        object_position = rand_xyz.squeeze()\n","        if object_type == 'block':\n","          object_shape = pybullet.createCollisionShape(pybullet.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])\n","          object_visual = pybullet.createVisualShape(pybullet.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])\n","          object_id = pybullet.createMultiBody(0.01, object_shape, object_visual, basePosition=object_position)\n","        elif object_type == 'bowl':\n","          object_position[2] = 0\n","          object_id = pybullet.loadURDF(\"bowl/bowl.urdf\", object_position, useFixedBase=1)\n","        pybullet.changeVisualShape(object_id, -1, rgbaColor=object_color)\n","        self.obj_name_to_id[obj_name] = object_id\n","    \n","\n","    # Re-enable rendering.\n","    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 1)\n","\n","    for _ in range(200):\n","      pybullet.stepSimulation()\n","\n","    # record object positions at reset\n","    self.init_pos = {name: self.get_obj_pos(name) for name in object_list}\n","\n","    return self.get_observation()\n","\n","  def servoj(self, joints):\n","    \"\"\"Move to target joint positions with position control.\"\"\"\n","    pybullet.setJointMotorControlArray(\n","      bodyIndex=self.robot_id,\n","      jointIndices=self.joint_ids,\n","      controlMode=pybullet.POSITION_CONTROL,\n","      targetPositions=joints,\n","      positionGains=[0.01]*6)\n","  \n","  def movep(self, position):\n","    \"\"\"Move to target end effector position.\"\"\"\n","    joints = pybullet.calculateInverseKinematics(\n","        bodyUniqueId=self.robot_id,\n","        endEffectorLinkIndex=self.tip_link_id,\n","        targetPosition=position,\n","        targetOrientation=pybullet.getQuaternionFromEuler(self.home_ee_euler),\n","        maxNumIterations=100)\n","    self.servoj(joints)\n","\n","  def get_ee_pos(self):\n","    ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])\n","    return ee_xyz\n","\n","  def step(self, action=None):\n","    \"\"\"Do pick and place motion primitive.\"\"\"\n","    pick_pos, place_pos = action['pick'].copy(), action['place'].copy()\n","\n","    # Set fixed primitive z-heights.\n","    hover_xyz = np.float32([pick_pos[0], pick_pos[1], 0.2])\n","    if pick_pos.shape[-1] == 2:\n","      pick_xyz = np.append(pick_pos, 0.025)\n","    else:\n","      pick_xyz = pick_pos\n","      pick_xyz[2] = 0.025\n","    if place_pos.shape[-1] == 2:\n","      place_xyz = np.append(place_pos, 0.15)\n","    else:\n","      place_xyz = place_pos\n","      place_xyz[2] = 0.15\n","\n","    # Move to object.\n","    ee_xyz = self.get_ee_pos()\n","    while np.linalg.norm(hover_xyz - ee_xyz) > 0.01:\n","      self.movep(hover_xyz)\n","      self.step_sim_and_render()\n","      ee_xyz = self.get_ee_pos()\n","\n","    while np.linalg.norm(pick_xyz - ee_xyz) > 0.01:\n","      self.movep(pick_xyz)\n","      self.step_sim_and_render()\n","      ee_xyz = self.get_ee_pos()\n","\n","    # Pick up object.\n","    self.gripper.activate()\n","    for _ in range(240):\n","      self.step_sim_and_render()\n","    while np.linalg.norm(hover_xyz - ee_xyz) > 0.01:\n","      self.movep(hover_xyz)\n","      self.step_sim_and_render()\n","      ee_xyz = self.get_ee_pos()\n","    \n","    for _ in range(50):\n","      self.step_sim_and_render()\n","    \n","    # Move to place location.\n","    while np.linalg.norm(place_xyz - ee_xyz) > 0.01:\n","      self.movep(place_xyz)\n","      self.step_sim_and_render()\n","      ee_xyz = self.get_ee_pos()\n","\n","    # Place down object.\n","    while (not self.gripper.detect_contact()) and (place_xyz[2] > 0.03):\n","      place_xyz[2] -= 0.001\n","      self.movep(place_xyz)\n","      for _ in range(3):\n","        self.step_sim_and_render()\n","    self.gripper.release()\n","    for _ in range(240):\n","      self.step_sim_and_render()\n","    place_xyz[2] = 0.2\n","    ee_xyz = self.get_ee_pos()\n","    while np.linalg.norm(place_xyz - ee_xyz) > 0.01:\n","      self.movep(place_xyz)\n","      self.step_sim_and_render()\n","      ee_xyz = self.get_ee_pos()\n","    place_xyz = np.float32([0, -0.5, 0.2])\n","    while np.linalg.norm(place_xyz - ee_xyz) > 0.01:\n","      self.movep(place_xyz)\n","      self.step_sim_and_render()\n","      ee_xyz = self.get_ee_pos()\n","\n","    observation = self.get_observation()\n","    reward = self.get_reward()\n","    done = False\n","    info = {}\n","    return observation, reward, done, info\n","\n","  def set_alpha_transparency(self, alpha: float) -> None:\n","    for id in range(20):\n","      visual_shape_data = pybullet.getVisualShapeData(id)\n","      for i in range(len(visual_shape_data)):\n","        object_id, link_index, _, _, _, _, _, rgba_color = visual_shape_data[i]\n","        rgba_color = list(rgba_color[0:3]) +  [alpha]\n","        pybullet.changeVisualShape(\n","            self.robot_id, linkIndex=i, rgbaColor=rgba_color)      \n","        pybullet.changeVisualShape(\n","            self.gripper.body, linkIndex=i, rgbaColor=rgba_color)\n","\n","  def step_sim_and_render(self):\n","    pybullet.stepSimulation()\n","    self.sim_step += 1\n","\n","    interval = 40 if self.high_frame_rate else 60\n","    # Render current image at 8 FPS.\n","    if self.sim_step % interval == 0 and self.render:\n","      self.cache_video.append(self.get_camera_image())\n","\n","  def get_camera_image(self):\n","    if not self.high_res:\n","      image_size = (240, 240)\n","      intrinsics = (120., 0, 120., 0, 120., 120., 0, 0, 1)\n","    else:\n","      image_size=(360, 360)\n","      intrinsics=(180., 0, 180., 0, 180., 180., 0, 0, 1)\n","    color, _, _, _, _ = env.render_image(image_size, intrinsics)\n","    return color\n","\n","  def get_reward(self):\n","    return None\n","\n","  def get_observation(self):\n","    observation = {}\n","\n","    # Render current image.\n","    color, depth, position, orientation, intrinsics = self.render_image()\n","\n","    # Get heightmaps and colormaps.\n","    points = self.get_pointcloud(depth, intrinsics)\n","    position = np.float32(position).reshape(3, 1)\n","    rotation = pybullet.getMatrixFromQuaternion(orientation)\n","    rotation = np.float32(rotation).reshape(3, 3)\n","    transform = np.eye(4)\n","    transform[:3, :] = np.hstack((rotation, position))\n","    points = self.transform_pointcloud(points, transform)\n","    heightmap, colormap, xyzmap = self.get_heightmap(points, color, BOUNDS, PIXEL_SIZE)\n","\n","    observation[\"image\"] = colormap\n","    observation[\"xyzmap\"] = xyzmap\n","\n","    return observation\n","\n","  def render_image(self, image_size=(720, 720), intrinsics=(360., 0, 360., 0, 360., 360., 0, 0, 1)):\n","\n","    # Camera parameters.\n","    position = (0, -0.85, 0.4)\n","    orientation = (np.pi / 4 + np.pi / 48, np.pi, np.pi)\n","    orientation = pybullet.getQuaternionFromEuler(orientation)\n","    zrange = (0.01, 10.)\n","    noise=True\n","\n","    # OpenGL camera settings.\n","    lookdir = np.float32([0, 0, 1]).reshape(3, 1)\n","    updir = np.float32([0, -1, 0]).reshape(3, 1)\n","    rotation = pybullet.getMatrixFromQuaternion(orientation)\n","    rotm = np.float32(rotation).reshape(3, 3)\n","    lookdir = (rotm @ lookdir).reshape(-1)\n","    updir = (rotm @ updir).reshape(-1)\n","    lookat = position + lookdir\n","    focal_len = intrinsics[0]\n","    znear, zfar = (0.01, 10.)\n","    viewm = pybullet.computeViewMatrix(position, lookat, updir)\n","    fovh = (image_size[0] / 2) / focal_len\n","    fovh = 180 * np.arctan(fovh) * 2 / np.pi\n","\n","    # Notes: 1) FOV is vertical FOV 2) aspect must be float\n","    aspect_ratio = image_size[1] / image_size[0]\n","    projm = pybullet.computeProjectionMatrixFOV(fovh, aspect_ratio, znear, zfar)\n","\n","    # Render with OpenGL camera settings.\n","    _, _, color, depth, segm = pybullet.getCameraImage(\n","        width=image_size[1],\n","        height=image_size[0],\n","        viewMatrix=viewm,\n","        projectionMatrix=projm,\n","        shadow=1,\n","        flags=pybullet.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,\n","        renderer=pybullet.ER_BULLET_HARDWARE_OPENGL)\n","\n","    # Get color image.\n","    color_image_size = (image_size[0], image_size[1], 4)\n","    color = np.array(color, dtype=np.uint8).reshape(color_image_size)\n","    color = color[:, :, :3]  # remove alpha channel\n","    if noise:\n","      color = np.int32(color)\n","      color += np.int32(np.random.normal(0, 3, color.shape))\n","      color = np.uint8(np.clip(color, 0, 255))\n","\n","    # Get depth image.\n","    depth_image_size = (image_size[0], image_size[1])\n","    zbuffer = np.float32(depth).reshape(depth_image_size)\n","    depth = (zfar + znear - (2 * zbuffer - 1) * (zfar - znear))\n","    depth = (2 * znear * zfar) / depth\n","    if noise:\n","      depth += np.random.normal(0, 0.003, depth.shape)\n","\n","    intrinsics = np.float32(intrinsics).reshape(3, 3)\n","    return color, depth, position, orientation, intrinsics\n","\n","  def get_pointcloud(self, depth, intrinsics):\n","    \"\"\"Get 3D pointcloud from perspective depth image.\n","    Args:\n","      depth: HxW float array of perspective depth in meters.\n","      intrinsics: 3x3 float array of camera intrinsics matrix.\n","    Returns:\n","      points: HxWx3 float array of 3D points in camera coordinates.\n","    \"\"\"\n","    height, width = depth.shape\n","    xlin = np.linspace(0, width - 1, width)\n","    ylin = np.linspace(0, height - 1, height)\n","    px, py = np.meshgrid(xlin, ylin)\n","    px = (px - intrinsics[0, 2]) * (depth / intrinsics[0, 0])\n","    py = (py - intrinsics[1, 2]) * (depth / intrinsics[1, 1])\n","    points = np.float32([px, py, depth]).transpose(1, 2, 0)\n","    return points\n","\n","  def transform_pointcloud(self, points, transform):\n","    \"\"\"Apply rigid transformation to 3D pointcloud.\n","    Args:\n","      points: HxWx3 float array of 3D points in camera coordinates.\n","      transform: 4x4 float array representing a rigid transformation matrix.\n","    Returns:\n","      points: HxWx3 float array of transformed 3D points.\n","    \"\"\"\n","    padding = ((0, 0), (0, 0), (0, 1))\n","    homogen_points = np.pad(points.copy(), padding,\n","                            'constant', constant_values=1)\n","    for i in range(3):\n","      points[Ellipsis, i] = np.sum(transform[i, :] * homogen_points, axis=-1)\n","    return points\n","\n","  def get_heightmap(self, points, colors, bounds, pixel_size):\n","    \"\"\"Get top-down (z-axis) orthographic heightmap image from 3D pointcloud.\n","    Args:\n","      points: HxWx3 float array of 3D points in world coordinates.\n","      colors: HxWx3 uint8 array of values in range 0-255 aligned with points.\n","      bounds: 3x2 float array of values (rows: X,Y,Z; columns: min,max) defining\n","        region in 3D space to generate heightmap in world coordinates.\n","      pixel_size: float defining size of each pixel in meters.\n","    Returns:\n","      heightmap: HxW float array of height (from lower z-bound) in meters.\n","      colormap: HxWx3 uint8 array of backprojected color aligned with heightmap.\n","      xyzmap: HxWx3 float array of XYZ points in world coordinates.\n","    \"\"\"\n","    width = int(np.round((bounds[0, 1] - bounds[0, 0]) / pixel_size))\n","    height = int(np.round((bounds[1, 1] - bounds[1, 0]) / pixel_size))\n","    heightmap = np.zeros((height, width), dtype=np.float32)\n","    colormap = np.zeros((height, width, colors.shape[-1]), dtype=np.uint8)\n","    xyzmap = np.zeros((height, width, 3), dtype=np.float32)\n","\n","    # Filter out 3D points that are outside of the predefined bounds.\n","    ix = (points[Ellipsis, 0] >= bounds[0, 0]) & (points[Ellipsis, 0] < bounds[0, 1])\n","    iy = (points[Ellipsis, 1] >= bounds[1, 0]) & (points[Ellipsis, 1] < bounds[1, 1])\n","    iz = (points[Ellipsis, 2] >= bounds[2, 0]) & (points[Ellipsis, 2] < bounds[2, 1])\n","    valid = ix & iy & iz\n","    points = points[valid]\n","    colors = colors[valid]\n","\n","    # Sort 3D points by z-value, which works with array assignment to simulate\n","    # z-buffering for rendering the heightmap image.\n","    iz = np.argsort(points[:, -1])\n","    points, colors = points[iz], colors[iz]\n","    px = np.int32(np.floor((points[:, 0] - bounds[0, 0]) / pixel_size))\n","    py = np.int32(np.floor((points[:, 1] - bounds[1, 0]) / pixel_size))\n","    px = np.clip(px, 0, width - 1)\n","    py = np.clip(py, 0, height - 1)\n","    heightmap[py, px] = points[:, 2] - bounds[2, 0]\n","    for c in range(colors.shape[-1]):\n","      colormap[py, px, c] = colors[:, c]\n","      xyzmap[py, px, c] = points[:, c]\n","    colormap = colormap[::-1, :, :]  # Flip up-down.\n","    xv, yv = np.meshgrid(np.linspace(BOUNDS[0, 0], BOUNDS[0, 1], height),\n","                         np.linspace(BOUNDS[1, 0], BOUNDS[1, 1], width))\n","    xyzmap[:, :, 0] = xv\n","    xyzmap[:, :, 1] = yv\n","    xyzmap = xyzmap[::-1, :, :]  # Flip up-down.\n","    heightmap = heightmap[::-1, :]  # Flip up-down.\n","    return heightmap, colormap, xyzmap\n","  \n","  def on_top_of(self, obj_a, obj_b):\n","    \"\"\"\n","    check if obj_a is on top of obj_b\n","    condition 1: l2 distance on xy plane is less than a threshold\n","    condition 2: obj_a is higher than obj_b\n","    \"\"\"\n","    obj_a_pos = self.get_obj_pos(obj_a)\n","    obj_b_pos = self.get_obj_pos(obj_b)\n","    xy_dist = np.linalg.norm(obj_a_pos[:2] - obj_b_pos[:2])\n","    if obj_b in CORNER_POS:\n","      is_near = xy_dist < 0.06\n","      return is_near\n","    elif 'bowl' in obj_b:\n","      is_near = xy_dist < 0.06\n","      is_higher = obj_a_pos[2] > obj_b_pos[2]\n","      return is_near and is_higher\n","    else:\n","      is_near = xy_dist < 0.04\n","      is_higher = obj_a_pos[2] > obj_b_pos[2]\n","      return is_near and is_higher\n","  \n","  def get_obj_id(self, obj_name):\n","    try:\n","      if obj_name in self.obj_name_to_id:\n","        obj_id = self.obj_name_to_id[obj_name]\n","      else:\n","        obj_name = obj_name.replace('circle', 'bowl').replace('square', 'block').replace('small', '').strip()\n","        obj_id = self.obj_name_to_id[obj_name]\n","    except:\n","      print(f'requested_name=\"{obj_name}\"')\n","      print(f'available_objects_and_id=\"{self.obj_name_to_id}')\n","    return obj_id\n","  \n","  def get_obj_pos(self, obj_name):\n","    obj_name = obj_name.replace('the', '').replace('_', ' ').strip()\n","    if obj_name in CORNER_POS:\n","      position = np.float32(np.array(CORNER_POS[obj_name]))\n","    else:\n","      pick_id = self.get_obj_id(obj_name)\n","      pose = pybullet.getBasePositionAndOrientation(pick_id)\n","      position = np.float32(pose[0])\n","    return position\n","  \n","  def get_bounding_box(self, obj_name):\n","    obj_id = self.get_obj_id(obj_name)\n","    return pybullet.getAABB(obj_id)"]},{"cell_type":"markdown","metadata":{"id":"XOqyIcuzoT3I"},"source":["## LMP Setup"]},{"cell_type":"markdown","metadata":{"id":"3mGyvT_rh7yh"},"source":["### LMP Wrapper"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"xMZCl27sh-sY"},"outputs":[],"source":["class LMP_wrapper():\n","\n","  def __init__(self, env, cfg, render=False):\n","    self.env = env\n","    self._cfg = cfg\n","    self.object_names = list(self._cfg['env']['init_objs'])\n","    \n","    self._min_xy = np.array(self._cfg['env']['coords']['bottom_left'])\n","    self._max_xy = np.array(self._cfg['env']['coords']['top_right'])\n","    self._range_xy = self._max_xy - self._min_xy\n","\n","    self._table_z = self._cfg['env']['coords']['table_z']\n","    self.render = render\n","\n","  def is_obj_visible(self, obj_name):\n","    return obj_name in self.object_names\n","\n","  def get_obj_names(self):\n","    return self.object_names[::]\n","\n","  def denormalize_xy(self, pos_normalized):\n","    return pos_normalized * self._range_xy + self._min_xy\n","\n","  def get_corner_positions(self):\n","    unit_square = box(0, 0, 1, 1)\n","    normalized_corners = np.array(list(unit_square.exterior.coords))[:4]\n","    corners = np.array(([self.denormalize_xy(corner) for corner in normalized_corners]))\n","    return corners\n","\n","  def get_side_positions(self):\n","    side_xs = np.array([0, 0.5, 0.5, 1])\n","    side_ys = np.array([0.5, 0, 1, 0.5])\n","    normalized_side_positions = np.c_[side_xs, side_ys]\n","    side_positions = np.array(([self.denormalize_xy(corner) for corner in normalized_side_positions]))\n","    return side_positions\n","\n","  def get_obj_pos(self, obj_name):\n","    # return the xy position of the object in robot base frame\n","    return self.env.get_obj_pos(obj_name)[:2]\n","\n","  def get_obj_position_np(self, obj_name):\n","    return self.get_pos(obj_name)\n","\n","  def get_bbox(self, obj_name):\n","    # return the axis-aligned object bounding box in robot base frame (not in pixels)\n","    # the format is (min_x, min_y, max_x, max_y)\n","    bbox = self.env.get_bounding_box(obj_name)\n","    return bbox\n","\n","  def get_color(self, obj_name):\n","    for color, rgb in COLORS.items():\n","      if color in obj_name:\n","        return rgb\n","\n","  def pick_place(self, pick_pos, place_pos):\n","    pick_pos_xyz = np.r_[pick_pos, [self._table_z]]\n","    place_pos_xyz = np.r_[place_pos, [self._table_z]]\n","    pass\n","\n","  def put_first_on_second(self, arg1, arg2):\n","    # put the object with obj_name on top of target\n","    # target can either be another object name, or it can be an x-y position in robot base frame\n","    pick_pos = self.get_obj_pos(arg1) if isinstance(arg1, str) else arg1\n","    place_pos = self.get_obj_pos(arg2) if isinstance(arg2, str) else arg2\n","    self.env.step(action={'pick': pick_pos, 'place': place_pos})\n","\n","  def get_robot_pos(self):\n","    # return robot end-effector xy position in robot base frame\n","    return self.env.get_ee_pos()\n","\n","  def goto_pos(self, position_xy):\n","    # move the robot end-effector to the desired xy position while maintaining same z\n","    ee_xyz = self.env.get_ee_pos()\n","    position_xyz = np.concatenate([position_xy, ee_xyz[-1]])\n","    while np.linalg.norm(position_xyz - ee_xyz) > 0.01:\n","      self.env.movep(position_xyz)\n","      self.env.step_sim_and_render()\n","      ee_xyz = self.env.get_ee_pos()\n","\n","  def follow_traj(self, traj):\n","    for pos in traj:\n","      self.goto_pos(pos)\n","\n","  def get_corner_positions(self):\n","    normalized_corners = np.array([\n","        [0, 1],\n","        [1, 1],\n","        [0, 0],\n","        [1, 0]\n","    ])\n","    return np.array(([self.denormalize_xy(corner) for corner in normalized_corners]))\n","\n","  def get_side_positions(self):\n","    normalized_sides = np.array([\n","        [0.5, 1],\n","        [1, 0.5],\n","        [0.5, 0],\n","        [0, 0.5]\n","    ])\n","    return np.array(([self.denormalize_xy(side) for side in normalized_sides]))\n","\n","  def get_corner_name(self, pos):\n","    corner_positions = self.get_corner_positions()\n","    corner_idx = np.argmin(np.linalg.norm(corner_positions - pos, axis=1))\n","    return ['top left corner', 'top right corner', 'bottom left corner', 'botom right corner'][corner_idx]\n","\n","  def get_side_name(self, pos):\n","    side_positions = self.get_side_positions()\n","    side_idx = np.argmin(np.linalg.norm(side_positions - pos, axis=1))\n","    return ['top side', 'right side', 'bottom side', 'left side'][side_idx]"]},{"cell_type":"markdown","metadata":{"id":"KgAAf2BNoWMJ"},"source":["### LMP Prompts"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"oCibEkErnp-K"},"outputs":[],"source":["prompt_tabletop_ui = '''\n","# Python 2D robot control script\n","import numpy as np\n","from env_utils import put_first_on_second, get_obj_pos, get_obj_names, say, get_corner_name, get_side_name, is_obj_visible, stack_objects_in_order\n","from plan_utils import parse_obj_name, parse_position, parse_question, transform_shape_pts\n","\n","objects = ['yellow block', 'green block', 'yellow bowl', 'blue block', 'blue bowl', 'green bowl']\n","# the yellow block on the yellow bowl.\n","say('Ok - putting the yellow block on the yellow bowl')\n","put_first_on_second('yellow block', 'yellow bowl')\n","objects = ['yellow block', 'green block', 'yellow bowl', 'blue block', 'blue bowl', 'green bowl']\n","# which block did you move.\n","say('I moved the yellow block')\n","objects = ['yellow block', 'green block', 'yellow bowl', 'blue block', 'blue bowl', 'green bowl']\n","# move the green block to the top right corner.\n","say('Got it - putting the green block on the top right corner')\n","corner_pos = parse_position('top right corner')\n","put_first_on_second('green block', corner_pos)\n","objects = ['yellow block', 'green block', 'yellow bowl', 'blue block', 'blue bowl', 'green bowl']\n","# stack the blue bowl on the yellow bowl on the green block.\n","order_bottom_to_top = ['green block', 'yellow block', 'blue bowl']\n","say(f'Sure - stacking from top to bottom: {\", \".join(order_bottom_to_top)}')\n","stack_objects_in_order(object_names=order_bottom_to_top)\n","objects = ['cyan block', 'white block', 'cyan bowl', 'blue block', 'blue bowl', 'white bowl']\n","# move the cyan block into its corresponding bowl.\n","matches = {'cyan block': 'cyan bowl'}\n","say('Got it - placing the cyan block on the cyan bowl')\n","for first, second in matches.items():\n","  put_first_on_second(first, get_obj_pos(second))\n","objects = ['cyan block', 'white block', 'cyan bowl', 'blue block', 'blue bowl', 'white bowl']\n","# make a line of blocks on the right side.\n","say('No problem! Making a line of blocks on the right side')\n","block_names = parse_obj_name('the blocks', f'objects = {get_obj_names()}')\n","line_pts = parse_position(f'a 30cm vertical line on the right with {len(block_names)} points')\n","for block_name, pt in zip(block_names, line_pts):\n","  put_first_on_second(block_name, pt)\n","objects = ['yellow block', 'red block', 'yellow bowl', 'gray block', 'gray bowl', 'red bowl']\n","# put the small banana colored thing in between the blue bowl and green block.\n","say('Sure thing - putting the yellow block between the blue bowl and the green block')\n","target_pos = parse_position('a point in the middle betweeen the blue bowl and the green block')\n","put_first_on_second('yellow block', target_pos)\n","objects = ['yellow block', 'red block', 'yellow bowl', 'gray block', 'gray bowl', 'red bowl']\n","# can you cut the bowls in half.\n","say('no, I can only move objects around')\n","objects = ['yellow block', 'green block', 'yellow bowl', 'gray block', 'gray bowl', 'green bowl']\n","# stack the blocks on the right side with the gray one on the bottom.\n","say('Ok. stacking the blocks on the right side with the gray block on the bottom')\n","right_side = parse_position('the right side')\n","put_first_on_second('gray block', right_side)\n","order_bottom_to_top = ['gray block', 'green block', 'yellow block']\n","stack_objects_in_order(object_names=order_bottom_to_top)\n","objects = ['yellow block', 'green block', 'yellow bowl', 'blue block', 'blue bowl', 'green bowl']\n","# hide the blue bowl.\n","bowl_name = np.random.choice(['yellow bowl', 'green bowl'])\n","say(f'Sounds good! Hiding the blue bowl under the {bowl_name}')\n","put_first_on_second(bowl_name, 'blue bowl')\n","objects = ['pink block', 'green block', 'pink bowl', 'blue block', 'blue bowl', 'green bowl']\n","# stack everything with the green block on top.\n","say('Ok! Stacking everything with the green block on the top')\n","order_bottom_to_top = ['blue bowl', 'pink bowl', 'green bowl', 'pink block', 'blue block', 'green block']\n","stack_objects_in_order(object_names=order_bottom_to_top)\n","objects = ['pink block', 'green block', 'pink bowl', 'blue block', 'blue bowl', 'green bowl']\n","# move the grass-colored bowl to the left.\n","say('Sure - moving the green bowl left by 10 centimeters')\n","left_pos = parse_position('a point 10cm left of the green bowl')\n","put_first_on_second('green bowl', left_pos)\n","objects = ['pink block', 'green block', 'pink bowl', 'blue block', 'blue bowl', 'green bowl']\n","# why did you move the red bowl.\n","say(f'I did not move the red bowl')\n","objects = ['pink block', 'green block', 'pink bowl', 'blue block', 'blue bowl', 'green bowl']\n","# undo that.\n","say('Sure - moving the green bowl right by 10 centimeters')\n","left_pos = parse_position('a point 10cm right of the green bowl')\n","put_first_on_second('green bowl', left_pos)\n","objects = ['brown bowl', 'green block', 'brown block', 'green bowl', 'blue bowl', 'blue block']\n","# place the top most block to the corner closest to the bottom most block.\n","top_block_name = parse_obj_name('top most block', f'objects = {get_obj_names()}')\n","bottom_block_name = parse_obj_name('bottom most block', f'objects = {get_obj_names()}')\n","closest_corner_pos = parse_position(f'the corner closest to the {bottom_block_name}', f'objects = {get_obj_names()}')\n","say(f'Putting the {top_block_name} on the {get_corner_name(closest_corner_pos)}')\n","put_first_on_second(top_block_name, closest_corner_pos)\n","objects = ['brown bowl', 'green block', 'brown block', 'green bowl', 'blue bowl', 'blue block']\n","# move the brown bowl to the side closest to the green block.\n","closest_side_position = parse_position('the side closest to the green block')\n","say(f'Got it - putting the brown bowl on the {get_side_name(closest_side_position)}')\n","put_first_on_second('brown bowl', closest_side_position)\n","objects = ['brown bowl', 'green block', 'brown block', 'green bowl', 'blue bowl', 'blue block']\n","# place the green block to the right of the bowl that has the blue block.\n","bowl_name = parse_obj_name('the bowl that has the blue block', f'objects = {get_obj_names()}')\n","if bowl_name:\n","  target_pos = parse_position(f'a point 10cm to the right of the {bowl_name}')\n","  say(f'No problem - placing the green block to the right of the {bowl_name}')\n","  put_first_on_second('green block', target_pos)\n","else:\n","  say('There are no bowls that has the blue block')\n","objects = ['brown bowl', 'green block', 'brown block', 'green bowl', 'blue bowl', 'blue block']\n","# place the blue block in the empty bowl.\n","empty_bowl_name = parse_obj_name('the empty bowl', f'objects = {get_obj_names()}')\n","if empty_bowl_name:\n","  say(f'Ok! Putting the blue block on the {empty_bowl_name}')\n","  put_first_on_second('blue block', empty_bowl_name)\n","else:\n","  say('There are no empty bowls')\n","objects = ['brown bowl', 'green block', 'brown block', 'green bowl', 'blue bowl', 'blue block']\n","# move the other blocks to the bottom corners.\n","block_names = parse_obj_name('blocks other than the blue block', f'objects = {get_obj_names()}')\n","corners = parse_position('the bottom corners')\n","for block_name, pos in zip(block_names, corners):\n","  put_first_on_second(block_name, pos)\n","objects = ['brown bowl', 'green block', 'brown block', 'green bowl', 'blue bowl', 'blue block']\n","# move the red bowl a lot to the left of the blocks.\n","say('Sure! Moving the red bowl to a point left of the blocks')\n","left_pos = parse_position('a point 20cm left of the blocks')\n","put_first_on_second('red bowl', left_pos)\n","objects = ['pink block', 'gray block', 'orange block']\n","# move the pinkish colored block on the bottom side.\n","say('Ok - putting the pink block on the bottom side')\n","bottom_side_pos = parse_position('the bottom side')\n","put_first_on_second('pink block', bottom_side_pos)\n","objects = ['yellow bowl', 'blue block', 'yellow block', 'blue bowl']\n","# is the blue block to the right of the yellow bowl?\n","if parse_question('is the blue block to the right of the yellow bowl?', f'objects = {get_obj_names()}'):\n","  say('yes, there is a blue block to the right of the yellow bow')\n","else:\n","  say('no, there is\\'t a blue block to the right of the yellow bow')\n","objects = ['yellow bowl', 'blue block', 'yellow block', 'blue bowl']\n","# how many yellow objects are there?\n","n_yellow_objs = parse_question('how many yellow objects are there', f'objects = {get_obj_names()}')\n","say(f'there are {n_yellow_objs} yellow object')\n","objects = ['pink block', 'green block', 'pink bowl', 'blue block', 'blue bowl', 'green bowl']\n","# move the left most block to the green bowl.\n","left_block_name = parse_obj_name('left most block', f'objects = {get_obj_names()}')\n","say(f'Moving the {left_block_name} on the green bowl')\n","put_first_on_second(left_block_name, 'green bowl')\n","objects = ['pink block', 'green block', 'pink bowl', 'blue block', 'blue bowl', 'green bowl']\n","# move the other blocks to different corners.\n","block_names = parse_obj_name(f'blocks other than the {left_block_name}', f'objects = {get_obj_names()}')\n","corners = parse_position('the corners')\n","say(f'Ok - moving the other {len(block_names)} blocks to different corners')\n","for block_name, pos in zip(block_names, corners):\n","  put_first_on_second(block_name, pos)\n","objects = ['pink block', 'green block', 'pink bowl', 'blue block', 'blue bowl', 'green bowl']\n","# is the pink block on the green bowl.\n","if parse_question('is the pink block on the green bowl', f'objects = {get_obj_names()}'):\n","  say('Yes - the pink block is on the green bowl.')\n","else:\n","  say('No - the pink block is not on the green bowl.')\n","objects = ['pink block', 'green block', 'pink bowl', 'blue block', 'blue bowl', 'green bowl']\n","# what are the blocks left of the green bowl.\n","left_block_names =  parse_question('what are the blocks left of the green bowl', f'objects = {get_obj_names()}')\n","if len(left_block_names) > 0:\n","  say(f'These blocks are left of the green bowl: {\", \".join(left_block_names)}')\n","else:\n","  say('There are no blocks left of the green bowl')\n","objects = ['pink block', 'green block', 'pink bowl', 'blue block', 'blue bowl', 'green bowl']\n","# if you see a purple bowl put it on the blue bowl\n","if is_obj_visible('purple bowl'):\n","  say('Putting the purple bowl on the pink bowl')\n","  put_first_on_second('purple bowl', 'pink bowl')\n","else:\n","  say('I don\\'t see a purple bowl')\n","objects = ['yellow block', 'green block', 'yellow bowl', 'blue block', 'blue bowl', 'green bowl']\n","# imagine that the bowls are different biomes on earth and imagine that the blocks are parts of a building.\n","say('ok')\n","objects = ['yellow block', 'green block', 'yellow bowl', 'blue block', 'blue bowl', 'green bowl']\n","# now build a tower in the grasslands.\n","order_bottom_to_top = ['green bowl', 'blue block', 'green block', 'yellow block']\n","say('stacking the blocks on the green bowl')\n","stack_objects_in_order(object_names=order_bottom_to_top)\n","objects = ['yellow block', 'green block', 'yellow bowl', 'gray block', 'gray bowl', 'green bowl']\n","# show me what happens when the desert gets flooded by the ocean.\n","say('putting the yellow bowl on the blue bowl')\n","put_first_on_second('yellow bowl', 'blue bowl')\n","objects = ['pink block', 'gray block', 'orange block']\n","# move all blocks 5cm toward the top.\n","say('Ok - moving all blocks 5cm toward the top')\n","block_names = parse_obj_name('the blocks', f'objects = {get_obj_names()}')\n","for block_name in block_names:\n","  target_pos = parse_position(f'a point 5cm above the {block_name}')\n","  put_first_on_second(block_name, target_pos)\n","objects = ['cyan block', 'white block', 'purple bowl', 'blue block', 'blue bowl', 'white bowl']\n","# make a triangle of blocks in the middle.\n","block_names = parse_obj_name('the blocks', f'objects = {get_obj_names()}')\n","triangle_pts = parse_position(f'a triangle with size 10cm around the middle with {len(block_names)} points')\n","say('Making a triangle of blocks around the middle of the workspace')\n","for block_name, pt in zip(block_names, triangle_pts):\n","  put_first_on_second(block_name, pt)\n","objects = ['cyan block', 'white block', 'purple bowl', 'blue block', 'blue bowl', 'white bowl']\n","# make the triangle smaller.\n","triangle_pts = transform_shape_pts('scale it by 0.5x', shape_pts=triangle_pts)\n","say('Making the triangle smaller')\n","block_names = parse_obj_name('the blocks', f'objects = {get_obj_names()}')\n","for block_name, pt in zip(block_names, triangle_pts):\n","  put_first_on_second(block_name, pt)\n","objects = ['brown bowl', 'red block', 'brown block', 'red bowl', 'pink bowl', 'pink block']\n","# put the red block on the farthest bowl.\n","farthest_bowl_name = parse_obj_name('the bowl farthest from the red block', f'objects = {get_obj_names()}')\n","say(f'Putting the red block on the {farthest_bowl_name}')\n","put_first_on_second('red block', farthest_bowl_name)\n","'''.strip()"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"28aPdrLRp7Z7"},"outputs":[],"source":["prompt_parse_obj_name = '''\n","import numpy as np\n","from env_utils import get_obj_pos, parse_position\n","from utils import get_obj_positions_np\n","\n","objects = ['blue block', 'cyan block', 'purple bowl', 'gray bowl', 'brown bowl', 'pink block', 'purple block']\n","# the block closest to the purple bowl.\n","block_names = ['blue block', 'cyan block', 'purple block']\n","block_positions = get_obj_positions_np(block_names)\n","closest_block_idx = get_closest_idx(points=block_positions, point=get_obj_pos('purple bowl'))\n","closest_block_name = block_names[closest_block_idx]\n","ret_val = closest_block_name\n","objects = ['brown bowl', 'banana', 'brown block', 'apple', 'blue bowl', 'blue block']\n","# the blocks.\n","ret_val = ['brown block', 'blue block']\n","objects = ['brown bowl', 'banana', 'brown block', 'apple', 'blue bowl', 'blue block']\n","# the brown objects.\n","ret_val = ['brown bowl', 'brown block']\n","objects = ['brown bowl', 'banana', 'brown block', 'apple', 'blue bowl', 'blue block']\n","# a fruit that's not the apple\n","fruit_names = ['banana', 'apple']\n","for fruit_name in fruit_names:\n","    if fruit_name != 'apple':\n","        ret_val = fruit_name\n","objects = ['blue block', 'cyan block', 'purple bowl', 'brown bowl', 'purple block']\n","# blocks above the brown bowl.\n","block_names = ['blue block', 'cyan block', 'purple block']\n","brown_bowl_pos = get_obj_pos('brown bowl')\n","use_block_names = []\n","for block_name in block_names:\n","    if get_obj_pos(block_name)[1] > brown_bowl_pos[1]:\n","        use_block_names.append(block_name)\n","ret_val = use_block_names\n","objects = ['blue block', 'cyan block', 'purple bowl', 'brown bowl', 'purple block']\n","# the blue block.\n","ret_val = 'blue block'\n","objects = ['blue block', 'cyan block', 'purple bowl', 'brown bowl', 'purple block']\n","# the block closest to the bottom right corner.\n","corner_pos = parse_position('bottom right corner')\n","block_names = ['blue block', 'cyan block', 'purple block']\n","block_positions = get_obj_positions_np(block_names)\n","closest_block_idx = get_closest_idx(points=block_positions, point=corner_pos)\n","closest_block_name = block_names[closest_block_idx]\n","ret_val = closest_block_name\n","objects = ['brown bowl', 'green block', 'brown block', 'green bowl', 'blue bowl', 'blue block']\n","# the left most block.\n","block_names = ['green block', 'brown block', 'blue block']\n","block_positions = get_obj_positions_np(block_names)\n","left_block_idx = np.argsort(block_positions[:, 0])[0]\n","left_block_name = block_names[left_block_idx]\n","ret_val = left_block_name\n","objects = ['brown bowl', 'green block', 'brown block', 'green bowl', 'blue bowl', 'blue block']\n","# the bowl on near the top.\n","bowl_names = ['brown bowl', 'green bowl', 'blue bowl']\n","bowl_positions = get_obj_positions_np(bowl_names)\n","top_bowl_idx = np.argsort(block_positions[:, 1])[-1]\n","top_bowl_name = bowl_names[top_bowl_idx]\n","ret_val = top_bowl_name\n","objects = ['yellow bowl', 'purple block', 'yellow block', 'purple bowl', 'pink bowl', 'pink block']\n","# the third bowl from the right.\n","bowl_names = ['yellow bowl', 'purple bowl', 'pink bowl']\n","bowl_positions = get_obj_positions_np(bowl_names)\n","bowl_idx = np.argsort(block_positions[:, 0])[-3]\n","bowl_name = bowl_names[bowl_idx]\n","ret_val = bowl_name\n","'''.strip()"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"vEgAqueLp_xy"},"outputs":[],"source":["prompt_parse_position = '''\n","import numpy as np\n","from shapely.geometry import *\n","from shapely.affinity import *\n","from env_utils import denormalize_xy, parse_obj_name, get_obj_names, get_obj_pos\n","\n","# a 30cm horizontal line in the middle with 3 points.\n","middle_pos = denormalize_xy([0.5, 0.5]) \n","start_pos = middle_pos + [-0.3/2, 0]\n","end_pos = middle_pos + [0.3/2, 0]\n","line = make_line(start=start_pos, end=end_pos)\n","points = interpolate_pts_on_line(line=line, n=3)\n","ret_val = points\n","# a 20cm vertical line near the right with 4 points.\n","middle_pos = denormalize_xy([1, 0.5]) \n","start_pos = middle_pos + [0, -0.2/2]\n","end_pos = middle_pos + [0, 0.2/2]\n","line = make_line(start=start_pos, end=end_pos)\n","points = interpolate_pts_on_line(line=line, n=4)\n","ret_val = points\n","# a diagonal line from the top left to the bottom right corner with 5 points.\n","top_left_corner = denormalize_xy([0, 1])\n","bottom_right_corner = denormalize_xy([1, 0])\n","line = make_line(start=top_left_corner, end=bottom_right_corner)\n","points = interpolate_pts_on_line(line=line, n=5)\n","ret_val = points\n","# a triangle with size 10cm with 3 points.\n","polygon = make_triangle(size=0.1, center=denormalize_xy([0.5, 0.5]))\n","points = get_points_from_polygon(polygon)\n","ret_val = points\n","# the corner closest to the sun colored block.\n","block_name = parse_obj_name('the sun colored block', f'objects = {get_obj_names()}')\n","corner_positions = np.array([denormalize_xy(pos) for pos in [[0, 0], [0, 1], [1, 1], [1, 0]]])\n","closest_corner_pos = get_closest_point(points=corner_positions, point=get_obj_pos(block_name))\n","ret_val = closest_corner_pos\n","# the side farthest from the right most bowl.\n","bowl_name = parse_obj_name('the right most bowl', f'objects = {get_obj_names()}')\n","side_positions = np.array([denormalize_xy(pos) for pos in [[0.5, 0], [0.5, 1], [1, 0.5], [0, 0.5]]])\n","farthest_side_pos = get_farthest_point(points=side_positions, point=get_obj_pos(bowl_name))\n","ret_val = farthest_side_pos\n","# a point above the third block from the bottom.\n","block_name = parse_obj_name('the third block from the bottom', f'objects = {get_obj_names()}')\n","ret_val = get_obj_pos(block_name) + [0.1, 0]\n","# a point 10cm left of the bowls.\n","bowl_names = parse_obj_name('the bowls', f'objects = {get_obj_names()}')\n","bowl_positions = get_all_object_positions_np(obj_names=bowl_names)\n","left_obj_pos = bowl_positions[np.argmin(bowl_positions[:, 0])] + [-0.1, 0]\n","ret_val = left_obj_pos\n","# the bottom side.\n","bottom_pos = denormalize_xy([0.5, 0])\n","ret_val = bottom_pos\n","# the top corners.\n","top_left_pos = denormalize_xy([0, 1])\n","top_right_pos = denormalize_xy([1, 1])\n","ret_val = [top_left_pos, top_right_pos]\n","'''.strip()"]},{"cell_type":"code","source":["prompt_parse_question = '''\n","from utils import get_obj_pos, get_obj_names, parse_obj_name, bbox_contains_pt, is_obj_visible\n","\n","objects = ['yellow bowl', 'blue block', 'yellow block', 'blue bowl', 'fruit', 'green block', 'black bowl']\n","# is the blue block to the right of the yellow bowl?\n","ret_val = get_obj_pos('blue block')[0] > get_obj_pos('yellow bowl')[0]\n","objects = ['yellow bowl', 'blue block', 'yellow block', 'blue bowl', 'fruit', 'green block', 'black bowl']\n","# how many yellow objects are there?\n","yellow_object_names = parse_obj_name('the yellow objects', f'objects = {get_obj_names()}')\n","ret_val = len(yellow_object_names)\n","objects = ['pink block', 'green block', 'pink bowl', 'blue block', 'blue bowl', 'green bowl']\n","# is the pink block on the green bowl?\n","ret_val = bbox_contains_pt(container_name='green bowl', obj_name='pink block')\n","objects = ['pink block', 'green block', 'pink bowl', 'blue block', 'blue bowl', 'green bowl']\n","# what are the blocks left of the green bowl?\n","block_names = parse_obj_name('the blocks', f'objects = {get_obj_names()}')\n","green_bowl_pos = get_obj_pos('green bowl')\n","left_block_names = []\n","for block_name in block_names:\n","  if get_obj_pos(block_name)[0] < green_bowl_pos[0]:\n","    left_block_names.append(block_name)\n","ret_val = left_block_names\n","objects = ['pink block', 'yellow block', 'pink bowl', 'blue block', 'blue bowl', 'yellow bowl']\n","# is the sun colored block above the blue bowl?\n","sun_block_name = parse_obj_name('sun colored block', f'objects = {get_obj_names()}')\n","sun_block_pos = get_obj_pos(sun_block_name)\n","blue_bowl_pos = get_obj_pos('blue bowl')\n","ret_val = sun_block_pos[1] > blue_bowl_pos[1]\n","objects = ['pink block', 'yellow block', 'pink bowl', 'blue block', 'blue bowl', 'yellow bowl']\n","# is the green block below the blue bowl?\n","ret_val = get_obj_pos('green block')[1] < get_obj_pos('blue bowl')[1]\n","'''.strip()"],"metadata":{"id":"h20VX6KSNAzL"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["prompt_transform_shape_pts = '''\n","import numpy as np\n","from utils import get_obj_pos, get_obj_names, parse_position, parse_obj_name\n","\n","# make it bigger by 1.5.\n","new_shape_pts = scale_pts_around_centroid_np(shape_pts, scale_x=1.5, scale_y=1.5)\n","# move it to the right by 10cm.\n","new_shape_pts = translate_pts_np(shape_pts, delta=[0.1, 0])\n","# move it to the top by 20cm.\n","new_shape_pts = translate_pts_np(shape_pts, delta=[0, 0.2])\n","# rotate it clockwise by 40 degrees.\n","new_shape_pts = rotate_pts_around_centroid_np(shape_pts, angle=-np.deg2rad(40))\n","# rotate by 30 degrees and make it slightly smaller\n","new_shape_pts = rotate_pts_around_centroid_np(shape_pts, angle=np.deg2rad(30))\n","new_shape_pts = scale_pts_around_centroid_np(new_shape_pts, scale_x=0.7, scale_y=0.7)\n","# move it toward the blue block.\n","block_name = parse_obj_name('the blue block', f'objects = {get_obj_names()}')\n","block_pos = get_obj_pos(block_name)\n","mean_delta = np.mean(block_pos - shape_pts, axis=1)\n","new_shape_pts = translate_pts_np(shape_pts, mean_delta)\n","'''.strip()"],"metadata":{"id":"9Ak3OR7fXiuY"},"execution_count":null,"outputs":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"BxEn8RsXqHup"},"outputs":[],"source":["prompt_fgen = '''\n","import numpy as np\n","from shapely.geometry import *\n","from shapely.affinity import *\n","\n","from env_utils import get_obj_pos, get_obj_names\n","from ctrl_utils import put_first_on_second\n","\n","# define function: total = get_total(xs=numbers).\n","def get_total(xs):\n","    return np.sum(xs)\n","\n","# define function: y = eval_line(x, slope, y_intercept=0).\n","def eval_line(x, slope, y_intercept):\n","    return x * slope + y_intercept\n","\n","# define function: pt = get_pt_to_the_left(pt, dist).\n","def get_pt_to_the_left(pt, dist):\n","    return pt + [-dist, 0]\n","\n","# define function: pt = get_pt_to_the_top(pt, dist).\n","def get_pt_to_the_top(pt, dist):\n","    return pt + [0, dist]\n","\n","# define function line = make_line_by_length(length=x).\n","def make_line_by_length(length):\n","  line = LineString([[0, 0], [length, 0]])\n","  return line\n","\n","# define function: line = make_vertical_line_by_length(length=x).\n","def make_vertical_line_by_length(length):\n","  line = make_line_by_length(length)\n","  vertical_line = rotate(line, 90)\n","  return vertical_line\n","\n","# define function: pt = interpolate_line(line, t=0.5).\n","def interpolate_line(line, t):\n","  pt = line.interpolate(t, normalized=True)\n","  return np.array(pt.coords[0])\n","\n","# example: scale a line by 2.\n","line = make_line_by_length(1)\n","new_shape = scale(line, xfact=2, yfact=2)\n","\n","# example: put object1 on top of object0.\n","put_first_on_second('object1', 'object0')\n","\n","# example: get the position of the first object.\n","obj_names = get_obj_names()\n","pos_2d = get_obj_pos(obj_names[0])\n","'''.strip()"]},{"cell_type":"markdown","source":["### LMP Config"],"metadata":{"id":"5NxXYsxWGL7L"}},{"cell_type":"code","execution_count":null,"metadata":{"id":"Ffz7T9C8oZPV"},"outputs":[],"source":["cfg_tabletop = {\n","  'lmps': {\n","    'tabletop_ui': {\n","      'prompt_text': prompt_tabletop_ui,\n","      'engine': model_name,\n","      'max_tokens': 512,\n","      'temperature': 0,\n","      'query_prefix': '# ',\n","      'query_suffix': '.',\n","      'stop': ['#', 'objects = ['],\n","      'maintain_session': True,\n","      'debug_mode': False,\n","      'include_context': True,\n","      'has_return': False,\n","      'return_val_name': 'ret_val',\n","    },\n","    'parse_obj_name': {\n","      'prompt_text': prompt_parse_obj_name,\n","      'engine': model_name,\n","      'max_tokens': 512,\n","      'temperature': 0,\n","      'query_prefix': '# ',\n","      'query_suffix': '.',\n","      'stop': ['#', 'objects = ['],\n","      'maintain_session': False,\n","      'debug_mode': False,\n","      'include_context': True,\n","      'has_return': True,\n","      'return_val_name': 'ret_val',\n","    },\n","    'parse_position': {\n","      'prompt_text': prompt_parse_position,\n","      'engine': model_name,\n","      'max_tokens': 512,\n","      'temperature': 0,\n","      'query_prefix': '# ',\n","      'query_suffix': '.',\n","      'stop': ['#'],\n","      'maintain_session': False,\n","      'debug_mode': False,\n","      'include_context': True,\n","      'has_return': True,\n","      'return_val_name': 'ret_val',\n","    },\n","    'parse_question': {\n","      'prompt_text': prompt_parse_question,\n","      'engine': model_name,\n","      'max_tokens': 512,\n","      'temperature': 0,\n","      'query_prefix': '# ',\n","      'query_suffix': '.',\n","      'stop': ['#', 'objects = ['],\n","      'maintain_session': False,\n","      'debug_mode': False,\n","      'include_context': True,\n","      'has_return': True,\n","      'return_val_name': 'ret_val',\n","    },\n","    'transform_shape_pts': {\n","      'prompt_text': prompt_transform_shape_pts,\n","      'engine': model_name,\n","      'max_tokens': 512,\n","      'temperature': 0,\n","      'query_prefix': '# ',\n","      'query_suffix': '.',\n","      'stop': ['#'],\n","      'maintain_session': False,\n","      'debug_mode': False,\n","      'include_context': True,\n","      'has_return': True,\n","      'return_val_name': 'new_shape_pts',\n","    },\n","    'fgen': {\n","      'prompt_text': prompt_fgen,\n","      'engine': model_name,\n","      'max_tokens': 512,\n","      'temperature': 0,\n","      'query_prefix': '# define function: ',\n","      'query_suffix': '.',\n","      'stop': ['# define', '# example'],\n","      'maintain_session': False,\n","      'debug_mode': False,\n","      'include_context': True,\n","    }\n","  }\n","}\n","\n","lmp_tabletop_coords = {\n","        'top_left':     (-0.3 + 0.05, -0.2 - 0.05),\n","        'top_side':     (0,           -0.2 - 0.05),\n","        'top_right':    (0.3 - 0.05,  -0.2 - 0.05),\n","        'left_side':    (-0.3 + 0.05, -0.5,      ),\n","        'middle':       (0,           -0.5,      ),\n","        'right_side':   (0.3 - 0.05,  -0.5,      ),\n","        'bottom_left':  (-0.3 + 0.05, -0.8 + 0.05),\n","        'bottom_side':  (0,           -0.8 + 0.05),\n","        'bottom_right': (0.3 - 0.05,  -0.8 + 0.05),\n","        'table_z':       0.0,\n","      }"]},{"cell_type":"markdown","metadata":{"id":"tP4a6iOi4rFn"},"source":["### LMP Utils"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"2spxf-mc4tOa"},"outputs":[],"source":["def setup_LMP(env, cfg_tabletop):\n","  # LMP env wrapper\n","  cfg_tabletop = copy.deepcopy(cfg_tabletop)\n","  cfg_tabletop['env'] = dict()\n","  cfg_tabletop['env']['init_objs'] = list(env.obj_name_to_id.keys())\n","  cfg_tabletop['env']['coords'] = lmp_tabletop_coords\n","  LMP_env = LMP_wrapper(env, cfg_tabletop)\n","  # creating APIs that the LMPs can interact with\n","  fixed_vars = {\n","      'np': np\n","  }\n","  fixed_vars.update({\n","      name: eval(name)\n","      for name in shapely.geometry.__all__ + shapely.affinity.__all__\n","  })\n","  variable_vars = {\n","      k: getattr(LMP_env, k)\n","      for k in [\n","          'get_bbox', 'get_obj_pos', 'get_color', 'is_obj_visible', 'denormalize_xy',\n","          'put_first_on_second', 'get_obj_names',\n","          'get_corner_name', 'get_side_name',\n","      ]\n","  }\n","  variable_vars['say'] = lambda msg: print(f'robot says: {msg}')\n","\n","  # creating the function-generating LMP\n","  lmp_fgen = LMPFGen(cfg_tabletop['lmps']['fgen'], fixed_vars, variable_vars)\n","\n","  # creating other low-level LMPs\n","  variable_vars.update({\n","      k: LMP(k, cfg_tabletop['lmps'][k], lmp_fgen, fixed_vars, variable_vars)\n","      for k in ['parse_obj_name', 'parse_position', 'parse_question', 'transform_shape_pts']\n","  })\n","\n","  # creating the LMP that deals w/ high-level language commands\n","  lmp_tabletop_ui = LMP(\n","      'tabletop_ui', cfg_tabletop['lmps']['tabletop_ui'], lmp_fgen, fixed_vars, variable_vars\n","  )\n","\n","  return lmp_tabletop_ui"]},{"cell_type":"markdown","metadata":{"id":"NYQsDt_LtlDk"},"source":["# Interactive Tabletop Manipulation\n","\n","Instructions: \n","\n","1. Set the number of blocks and bowls with the sliders in the next cell.\n","2. Input a command in the `user_input` field below and run the cell. This will run Code as Policies by querying the code-generating LLM to write robot code to complete the task.\n","\n","Note the object names printed after the Initilize Env cell - these are the objects in the scene and can be referred to in the commands.\n","\n","Supported commands:\n","* Spatial reasoning (e.g. to the left of the red block, the closest corner, the farthest bowl, the second block from the right)\n","* Sequential actions (e.g. put blocks in matching bowls, stack blocks on the bottom right corner)\n","* Contextual commands (e.g. do the same with the blue block, undo that)\n","* Language-based reasoning (e.g. put the forest-colored block on the ocean-colored bowl).\n","* Simple Q&A (e.g. how many blocks are to the left of the blue bowl?)\n","\n","Example commands (note object names may need to be changed depending the sampled object names):\n","* put the sun-colored block on the bowl closest to it\n","* stack the blocks on the bottom most bowl\n","* arrange the blocks as a square in the middle\n","* move the square 5cm to the right\n","* how many blocks are to the right of the orange bowl?\n","* pick up the block closest to the top left corner and place it on the bottom right corner\n","\n","Known limitations:\n","* In simulation we're using ground truth object poses instead of using vision models. This means that commands the require knowledge of visual apperances (e.g. darkest bowl, largest object) are not supported.\n","* Currently, the low-level pick place primitive does not do collision checking, so if there are many objects on the table, placing actions may incur collisions.\n","* Prompt saturation - if too many commands (10+) are executed in a row, then the LLM may start to ignore examples in the early parts of the prompt.\n","* Ambiguous instructions - if a given instruction doesn't lead to the desired actions, try rephrasing it to remove ambiguities (e.g. place the block on the closest bowl -> place the block on its closest bowl)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"FwkGynMZKnBs","cellView":"form"},"outputs":[],"source":["#@title Initialize Env { vertical-output: true }\n","num_blocks = 3 #@param {type:\"slider\", min:0, max:4, step:1}\n","num_bowls = 3 #@param {type:\"slider\", min:0, max:4, step:1}\n","high_resolution = False #@param {type:\"boolean\"}\n","high_frame_rate = False #@param {type:\"boolean\"}\n","\n","# setup env and LMP\n","env = PickPlaceEnv(render=True, high_res=high_resolution, high_frame_rate=high_frame_rate)\n","block_list = np.random.choice(ALL_BLOCKS, size=num_blocks, replace=False).tolist()\n","bowl_list = np.random.choice(ALL_BOWLS, size=num_bowls, replace=False).tolist()\n","obj_list = block_list + bowl_list\n","_ = env.reset(obj_list)\n","lmp_tabletop_ui = setup_LMP(env, cfg_tabletop)\n","\n","# display env\n","cv2_imshow(cv2.cvtColor(env.get_camera_image(), cv2.COLOR_BGR2RGB))\n","\n","print('available objects:')\n","print(obj_list)"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"bKa2_whYMa_U","cellView":"form"},"outputs":[],"source":["#@title Interactive Demo { vertical-output: true }\n","\n","user_input = 'put the blue block on the yellow bowl' #@param {allow-input: true, type:\"string\"}\n","\n","env.cache_video = []\n","\n","print('Running policy and recording video...')\n","lmp_tabletop_ui(user_input, f'objects = {env.object_list}')\n","\n","# render video\n","if env.cache_video:\n","  rendered_clip = ImageSequenceClip(env.cache_video, fps=35 if high_frame_rate else 25)\n","  display(rendered_clip.ipython_display(autoplay=1, loop=1))"]}],"metadata":{"colab":{"collapsed_sections":[],"provenance":[{"file_id":"1C53APdZO8IfMw3tlTQiC8_xuxhb8LFJt","timestamp":1663291616525},{"file_id":"11-tchi4DhfmCFk4fE_p2rlrmAYaRcOkf","timestamp":1663261808829}]},"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"name":"python"}},"nbformat":4,"nbformat_minor":0}